// #include <fstream>
#include <plan_manage/planner_manager.h>
#include <thread>
#include "visualization_msgs/Marker.h" // zx-todo

namespace ego_planner {

PlanManager::PlanManager() {}

PlanManager::~PlanManager() { std::cout << "des manager" << std::endl; }

void PlanManager::initPlanModules(ros::NodeHandle &nh, Visualization::Ptr vis)
{
  /* read algorithm parameters */
  nh.param("manager/max_vel", pp_.max_vel_, -1.0);
  nh.param("manager/max_acc", pp_.max_acc_, -1.0);
  nh.param("manager/max_jerk", pp_.max_jerk_, -1.0);
  nh.param("manager/feasibility_tolerance", pp_.feasibility_tolerance_, 0.0);
  nh.param("manager/control_points_distance", pp_.ctrl_pt_dist, -1.0);
  nh.param("manager/planning_horizon", pp_.planning_horizen_, 5.0);
  nh.param("manager/use_distinctive_trajs", pp_.use_distinctive_trajs, false);
  nh.param("manager/drone_id", pp_.drone_id, -1);

  nh.param("manager_yaw/gaze_time", py_.gaze_time, 0.2);
  nh.param("manager_yaw/forward_time", py_.forward_time, 0.2);
  nh.param("manager_yaw/interval_time", py_.interval_time, 0.5);
  nh.param("manager_yaw/max_yaw_rate", py_.max_yaw_rate, M_PI);
  nh.param("manager_yaw/key_yaw_gain", py_.key_yaw_gain, 2.0);
  nh.param("manager_yaw/max_sample_num", py_.max_sample_num, 6);

  py_.runtimes_ratio = 1 - py_.gaze_time / py_.interval_time;

  /* init variable */
  grid_map_.reset(new EntropyMap);
  grid_map_->initMap(nh);

  obs_poses_ = std::make_shared<std::vector<ObsPosPtr>>();
  obs_trajs_ = std::make_shared<std::vector<ObsTrajPtr>>();

  bspline_optimizer_.reset(new BsplineOptimizer);
  bspline_optimizer_->setParam(nh);
  bspline_optimizer_->setEnvironment(grid_map_);
  bspline_optimizer_->Astar_.reset(new AStar);
  bspline_optimizer_->Astar_->initGridMap(grid_map_, Eigen::Vector3i(100, 100, 40)); // TODO 这一句可以加到setEnvironment
  bspline_optimizer_->setObsPoses(obs_poses_);
  bspline_optimizer_->setObsTrajs(obs_trajs_);

  visualization_ = vis;

  /* init predictor */
  predictor_.reset(new Predictor);
  predictor_->init(nh, visualization_, obs_poses_, obs_trajs_);

  // if (pp_.drone_id == 0)
  // {
  //   kino_path_finder_.reset(new KinodynamicAstar(nh, grid_map_));
  // }

  /* init yaw planner */
  gimOpter_.setLimits(py_.max_yaw_rate, py_.gaze_time);
}

void PlanManager::angleLimite(double &angle)
{
  while (angle > M_PI)
  {
    angle -= (M_PI * 2);
  }
  while (angle < -M_PI)
  {
    angle += (M_PI * 2);
  }
}

// TODO: need to be deleted
bool PlanManager::kinodynamicReplan(Eigen::Vector3d start_pt,
                                    Eigen::Vector3d start_vel,
                                    Eigen::Vector3d start_acc,
                                    Eigen::Vector3d end_pt,
                                    Eigen::Vector3d end_vel)
{
  if ((start_pt - end_pt).norm() < 0.2)
  {
    cout << "[kino replan]: Close goal" << endl;
    return false;
  }

  // kinodynamic path searching
  Eigen::Matrix3d iniState;
  iniState.col(0) = start_pt;
  iniState.col(1) = start_vel;
  iniState.col(2) = start_acc;
  int id = 1; // target drone id
  double time_offset = ros::Time::now().toSec() - swarm_trajs_buf_.at(id).start_time_.toSec();
  if (!kino_path_finder_->search(swarm_trajs_buf_.at(id).traj_pos_, iniState, time_offset))
  {
    return false;
  }

  // parameterize the path to bspline
  double ts = kino_path_finder_->getT() / 20;
  vector<Eigen::Vector3d> points, terminal_constraint;
  kino_path_finder_->getSamples(ts, points, terminal_constraint);

  Eigen::MatrixXd ctrl_pts;
  UniformBspline::BsplineFitting(ts, points, terminal_constraint, ctrl_pts);
  visualization_->displayOptimalList(ctrl_pts, 1);

  // bspline init
  bspline_optimizer_->setPosControlPointsAndTs(ctrl_pts, ts);

  UniformBspline pos = UniformBspline(ctrl_pts, 3, ts);
  updateTrajInfo(pos);

  return true;
}

/**
 * @brief main func
 *
 * @param flag_polyInit whether init by min-snap poly traj
 * @param flag_randomPolyTraj if init by poly, whether random generate
 */
bool PlanManager::reboundReplan(Eigen::Vector3d start_p,
                                Eigen::Vector3d start_v,
                                Eigen::Vector3d start_a,
                                Eigen::Vector3d end_p,
                                Eigen::Vector3d end_v,
                                bool flag_polyInit,
                                bool flag_randomPolyTraj)
{
  printf("\033[47;30m\n[drone %d reboundReplan %d]==============================\033[0m\n",
         pp_.drone_id, replan_cnt_++);

  double dist = (start_p - end_p).norm();
  if (dist < 0.2)
  {
    cout << "[EGORP]: Close to goal" << endl;
    continous_failures_count_++;
    return false;
  }

  bspline_optimizer_->setLocalTargetPt(end_p);

  ros::Time t_start = ros::Time::now();
  ros::Duration t_init, t_opt, t_refine;

  /*** STEP 1: INIT ***/
  // pp_.ctrl_pt_dist/pp_.max_vel_ is too tense, and will surely exceed the acc/vel limits
  double deltaT = pp_.ctrl_pt_dist / pp_.max_vel_ * 1.2;
  vector<Eigen::Vector3d> pts_init, terminal_constraint;
  bool flag_regenerate = false;
  do
  {
    pts_init.clear();
    terminal_constraint.clear();
    flag_regenerate = false;

    /* Initial path generated from a min-snap traj by order. */
    // TODO 这里加一个，如果局部轨迹是空的话
    if (flag_first_call_ || flag_polyInit || flag_force_polynomial_)
    {
      flag_first_call_ = false;
      flag_force_polynomial_ = false;

      PolynomialTraj gl_traj;

      // 先加速后减速
      double dis_acc = pow(pp_.max_vel_, 2) / pp_.max_acc_;
      double time = dis_acc > dist
                        ? sqrt(dist / pp_.max_acc_)
                        : (dist - dis_acc) / pp_.max_vel_ +
                              2 * pp_.max_vel_ / pp_.max_acc_;

      if (!flag_randomPolyTraj)
      {
        gl_traj =
            PolynomialTraj::one_seg_traj_gen(start_p, start_v, start_a,
                                             end_p, end_v,
                                             Eigen::Vector3d::Zero(), time);
      }
      else
      {
        Eigen::Vector3d dir_h = ((start_p - end_p).cross(Eigen::Vector3d(0, 0, 1)))
                                    .normalized();
        Eigen::Vector3d dir_v = ((start_p - end_p).cross(dir_h))
                                    .normalized();
        // TODO MZY 0.978 and 0.989 are what?
        double coef = 0.4 * (-0.978 / (continous_failures_count_ + 0.989) + 0.989);
        double rand_h = ((double)rand()) / RAND_MAX - 0.5;
        double rand_v = ((double)rand()) / RAND_MAX - 0.5;
        Eigen::Vector3d rand_pt = (start_p + end_p) / 2 +
                                  rand_h * dist * dir_h * coef * 2 +
                                  rand_v * dist * dir_v * coef;
        Eigen::MatrixXd pos(3, 3);
        pos.col(0) = start_p;
        pos.col(1) = rand_pt;
        pos.col(2) = end_p;
        Eigen::VectorXd t(2);
        t(0) = t(1) = time / 2;
        gl_traj = PolynomialTraj::minSnapTraj(pos, start_v, end_v, start_a,
                                              Eigen::Vector3d::Zero(), t);
      }

      double t;
      bool flag_too_far;
      deltaT *= 1.5; // deltaT will be divided by 1.5 in the next
      do
      {
        deltaT /= 1.5;
        pts_init.clear();
        flag_too_far = false;
        Eigen::Vector3d last_pt = gl_traj.getPos(0);
        for (t = 0; t < time; t += deltaT)
        {
          Eigen::Vector3d pt = gl_traj.getPos(t);
          if ((last_pt - pt).norm() > pp_.ctrl_pt_dist * 1.5)
          {
            flag_too_far = true;
            break;
          }
          last_pt = pt;
          pts_init.push_back(pt);
        }
      } while (flag_too_far || pts_init.size() < 7);
      // ^^^ To make sure the initial path has enough points
      //     and the dis between ponits are not long

      t -= deltaT;
      terminal_constraint.push_back(gl_traj.getVel(0));
      terminal_constraint.push_back(end_v);
      terminal_constraint.push_back(gl_traj.getAcc(0));
      terminal_constraint.push_back(gl_traj.getAcc(t));
    }
    /* Initial path generated from previous trajectory. */
    else
    {
      double t;
      double t_cur = (ros::Time::now() - local_data_.start_time_).toSec();

      /* record the trajectory length up to each node */
      vector<double> arc_length;
      arc_length.push_back(0.0);

      /**
       * Store the points from previous traj
       * and new traj generated by min-snap if needed
       */
      vector<Eigen::Vector3d> pts;
      for (t = t_cur; t < local_data_.duration_ + 1e-3; t += deltaT)
      {
        pts.push_back(local_data_.traj_pos_.getPosT(t));
        if (t > t_cur)
        {
          arc_length.push_back(
              (pts.back() - pts[pts.size() - 2]).norm() + arc_length.back());
        }
      }

      /* If there is a distance between A and B, connect them with a polynomial */
      t -= deltaT;
      double poly_time = (local_data_.traj_pos_.getPosT(t) - end_p).norm() /
                         pp_.max_vel_ * 2;
      if (poly_time > deltaT)
      {
        PolynomialTraj gl_traj =
            PolynomialTraj::one_seg_traj_gen(local_data_.traj_pos_.getPosT(t),
                                             local_data_.traj_vel_.getPosT(t),
                                             local_data_.traj_acc_.getPosT(t),
                                             end_p, end_v, Eigen::Vector3d::Zero(),
                                             poly_time);
        for (t = deltaT; t < poly_time; t += deltaT)
        {
          if (!arc_length.empty())
          {
            pts.push_back(gl_traj.getPos(t));
            arc_length.push_back(
                (pts.back() - pts[pts.size() - 2]).norm() + arc_length.back());
          }
          else
          {
            ROS_ERROR("arc_length is empty, return!, why?");
            continous_failures_count_++;
            return false;
          }
        }
      }

      /**
       * Linearly interpolate sufficient points with equal spacing
       * on the previously obtained points.
       */
      double sampleL = 0;
      double cps_dist = pp_.ctrl_pt_dist * 1.5; // be divided by 1.5 at beginning
      size_t id = 0;
      do
      {
        cps_dist /= 1.5;
        pts_init.clear();
        sampleL = 0;
        id = 0;
        while ((id <= arc_length.size() - 2) && sampleL <= arc_length.back())
        {
          double arcL = arc_length[id];
          double arcR = arc_length[id + 1];
          if (sampleL >= arcL && sampleL < arcR)
          {
            // interpolation
            pts_init.push_back((sampleL - arcL) / (arcR - arcL) * pts[id + 1] +
                               (arcR - sampleL) / (arcR - arcL) * pts[id]);
            sampleL += cps_dist;
          }
          else
            id++;
        }
        pts_init.push_back(end_p);
      } while (pts_init.size() < 7);
      // ^^^ If the start point is very close to end point, this will help

      terminal_constraint.push_back(local_data_.traj_vel_.getPosT(t_cur));
      terminal_constraint.push_back(end_v);
      terminal_constraint.push_back(local_data_.traj_acc_.getPosT(t_cur));
      terminal_constraint.push_back(Eigen::Vector3d::Zero());

      // The initial path is unnormally too long!
      if (pts_init.size() > pp_.planning_horizen_ / pp_.ctrl_pt_dist * 3)
      {
        flag_force_polynomial_ = true;
        flag_regenerate = true;
      }
    }
  } while (flag_regenerate);

  Eigen::MatrixXd pts_ctrl, pts_ctrl_temp;
  UniformBspline::BsplineFitting(deltaT, pts_init, terminal_constraint, pts_ctrl);

  // TODO 返回值是bool;
  bspline_optimizer_->InitAndProcess(pts_ctrl);

  t_init = ros::Time::now() - t_start;
  t_start = ros::Time::now();

  /*** STEP 2: OPTIMIZE ***/
  bool flag_step_1_success = false;
  vector<vector<Eigen::Vector3d>> vis_trajs;

  if (pp_.use_distinctive_trajs)
  {
    std::vector<ControlPoints> trajs = bspline_optimizer_->distinctiveTrajs();
    // cout << "\033[1;33m[OPT]: # multi-trajs = " << trajs.size() << "\033[1;0m" << endl;

    double cost;
    double cost_min = std::numeric_limits<double>::max();
    for (int i = trajs.size() - 1; i >= 0; i--)
    {
      if (bspline_optimizer_->OptimizeTraj(pts_ctrl_temp, cost, trajs[i], deltaT))
      {
        // cout << "[OPT]: traj " << trajs.size() - i << " success." << endl;
        flag_step_1_success = true;
        if (cost < cost_min)
        {
          cost_min = cost;
          pts_ctrl = pts_ctrl_temp;
        }

        // visualization
        pts_init.clear();
        for (int j = 0; j < pts_ctrl_temp.cols(); j++) {
          pts_init.push_back(pts_ctrl_temp.col(j));
        }
        vis_trajs.push_back(pts_init);
      }
      else {
        // cout << "[OPT]: traj " << trajs.size() - i << " failed." << endl;
      }
    }

    t_opt = ros::Time::now() - t_start;

    // This visuallization will take up several milliseconds.
    visualization_->displayMultiInitPathList(vis_trajs, 0.2);
  }
  else
  {
    flag_step_1_success = bspline_optimizer_->OptimizeTraj(pts_ctrl, deltaT);
    t_opt = ros::Time::now() - t_start;
    pts_init.clear();
    for (int j = 0; j < pts_ctrl_temp.cols(); j++) {
      pts_init.push_back(pts_ctrl_temp.col(j));
    }
    visualization_->displayInitPathList(pts_init, 0.2, 0);
  }

  if (!flag_step_1_success)
  {
    visualization_->displayOptimalList(pts_ctrl, 0);
    continous_failures_count_++;
    return false;
  }

  t_start = ros::Time::now();

  /*** STEP 3: REFINE(RE-ALLOCATE TIME) IF NECESSARY ***/
  // TODO refine的部分我现在感觉是有大问题的，他在costfun里面也没有加上障碍物的惩罚项
  // 现在先不管了
  UniformBspline pos = UniformBspline(pts_ctrl, 3, deltaT);
  pos.setPhysicalLimits(pp_.max_vel_, pp_.max_acc_, pp_.feasibility_tolerance_);

  double ratio;
  bool flag_step_2_success = true;
  if (!pos.checkFeasibility(ratio, false))
  {
    cout << "u_.rows() =" << endl;
    cout << "[OPT]: Need to reallocate time." << endl;
    cout << "[OPT]: Need to reallocate time." << endl;

    Eigen::MatrixXd optimal_control_points;
    flag_step_2_success = refineTraj(pos, terminal_constraint, ratio, deltaT, optimal_control_points);
    if (flag_step_2_success)
      pos = UniformBspline(optimal_control_points, 3, deltaT);
  }

  if (!flag_step_2_success)
  {
    cout << "\033[34m[OPT]: This refined trajectory hits obstacles. "
         << "It doesn't matter if appeares occasionally. "
         << "But if continously appearing, Increase parameter \"lambda_fitness\".\033[0m" << endl;
    continous_failures_count_++;
    return false;
  }

  t_refine = ros::Time::now() - t_start;

  // Save planned results
  updateTrajInfo(pos);

  static double sum_time = 0;
  static int count_success = 0;
  sum_time += (t_init + t_opt + t_refine).toSec();
  count_success++;
  // cout << "[OPT]: Total time:\033[32m" << (t_init + t_opt + t_refine).toSec() << "\033[0m, optimize:"
  //      << (t_init + t_opt).toSec() << ", refine:" << t_refine.toSec() << ", avg_time=" << sum_time / count_success << endl;

  // success. YoY
  continous_failures_count_ = 0;
  return true;
}

bool PlanManager::reboundReplanForPredict(Eigen::MatrixXd &ctrl_pts, double ts)
{
  printf("\033[47;30m\n[The predict is in collision! Do a easy planning!]\033[0m\n");

  bspline_optimizer_->InitAndProcess(ctrl_pts);

  /*** STEP 2: OPTIMIZE ***/
  bool plan_success = bspline_optimizer_->OptimizeTraj(ctrl_pts, ts);

  return plan_success;
}

bool PlanManager::EmergencyStop(Eigen::Vector3d stop_pos)
{
  Eigen::MatrixXd control_points(3, 6);
  for (int i = 0; i < 6; i++)
  {
    control_points.col(i) = stop_pos;
  }

  updateTrajInfo(UniformBspline(control_points, 3, 1.0));

  return true;
}

bool PlanManager::checkCollision(int drone_id)
{
  // It means my first planning has not started
  if (local_data_.start_time_ == ros::Time(0))
    return false;

  double t_self_start = local_data_.start_time_.toSec();
  double t_other_start = swarm_trajs_buf_[drone_id].start_time_.toSec();
  double t_now = ros::Time::now().toSec();

  double t_start = max(max(t_self_start, t_other_start), t_now);
  double t_end = min(t_self_start + local_data_.duration_ * 2 / 3,
                     t_other_start + swarm_trajs_buf_[drone_id].duration_);

  UniformBspline other_traj = swarm_trajs_buf_[drone_id].traj_pos_;
  for (double t = t_start; t < t_end; t += 0.03)
  {
    if ((local_data_.traj_pos_.getPosT(t - t_self_start) -
         other_traj.getPosT(t - t_other_start))
            .norm() < bspline_optimizer_->getSwarmClearance())
    {
      return true;
    }
  }

  return false;
}

/**
 * @brief Generate global reference trajectory by Minimun-Snap
 */
bool PlanManager::planGlobalTraj(Eigen::Vector3d start_pos, Eigen::Vector3d start_vel,
                                 Eigen::Vector3d start_acc, Eigen::Vector3d end_pos,
                                 Eigen::Vector3d end_vel, Eigen::Vector3d end_acc)
{
  vector<Eigen::Vector3d> points;
  points.push_back(start_pos);
  points.push_back(end_pos);

  /* insert intermediate points if too far */
  vector<Eigen::Vector3d> inter_points;
  const double dist_thresh = 4.0;

  for (size_t i = 0; i < points.size() - 1; ++i)
  {
    inter_points.push_back(points.at(i));
    double dist = (points.at(i + 1) - points.at(i)).norm();
    if (dist > dist_thresh)
    {
      int id_num = floor(dist / dist_thresh) + 1;

      for (int j = 1; j < id_num; ++j)
      {
        Eigen::Vector3d inter_pt = points.at(i) * (1.0 - double(j) / id_num) +
                                   points.at(i + 1) * double(j) / id_num;
        inter_points.push_back(inter_pt);
      }
    }
  }
  inter_points.push_back(points.back());

  /* write position matrix and time vector */
  size_t pt_num = inter_points.size();
  Eigen::MatrixXd pos(3, pt_num);
  for (size_t i = 0; i < pt_num; ++i)
    pos.col(i) = inter_points[i];

  Eigen::VectorXd time(pt_num - 1);
  for (size_t i = 0; i < pt_num - 1; ++i)
  {
    time(i) = (pos.col(i + 1) - pos.col(i)).norm() / (pp_.max_vel_);
  }
  time(0) *= 2.0;
  time(time.rows() - 1) *= 2.0;

  PolynomialTraj gl_traj;
  if (pos.cols() >= 3)
    gl_traj = PolynomialTraj::minSnapTraj(pos, start_vel, end_vel,
                                          start_acc, end_acc, time);
  else if (pos.cols() == 2)
    gl_traj = PolynomialTraj::one_seg_traj_gen(start_pos, start_vel, start_acc,
                                               end_pos, end_vel, end_acc, time(0));
  else
    return false;

  global_data_.setGlobalTraj(gl_traj, ros::Time::now());

  return true;
}

bool PlanManager::refineTraj(UniformBspline &traj,
                             vector<Eigen::Vector3d> &terminal_constraint,
                             double ratio,
                             double &ts,
                             Eigen::MatrixXd &optimal_control_points)
{
  Eigen::MatrixXd ctrl_pts;
  reparamBspline(traj, terminal_constraint, ratio, ctrl_pts, ts);

  traj = UniformBspline(ctrl_pts, 3, ts);

  double t_step = traj.getTimeSum() / (ctrl_pts.cols() - 3);
  bspline_optimizer_->ref_pts_.clear();
  for (double t = 0; t < traj.getTimeSum() + 1e-4; t += t_step)
    bspline_optimizer_->ref_pts_.push_back(traj.getPosT(t));

  bool success = bspline_optimizer_->OptimizeTrajRefine(ctrl_pts, ts, optimal_control_points);

  return success;
}

void PlanManager::updateTrajInfo(const UniformBspline &position_traj)
{
  local_data_.start_time_ = ros::Time::now();
  local_data_.traj_pos_ = position_traj;
  local_data_.traj_vel_ = local_data_.traj_pos_.getDerivative();
  local_data_.traj_acc_ = local_data_.traj_vel_.getDerivative();
  local_data_.start_pos_ = local_data_.traj_pos_.getPosT(0.0);
  local_data_.duration_ = local_data_.traj_pos_.getTimeSum();
  local_data_.traj_id_ += 1;
}

void PlanManager::updateYawTrajInfo(const UniformBspline &yaw_traj, const ros::Time time_now)
{
  local_data_.start_time_ = time_now;
  local_data_.yaw_traj_ = yaw_traj;
  local_data_.yawdot_traj_ = local_data_.yaw_traj_.getDerivative();
  // local_data_.start_yaw_ = local_data_.yaw_traj.getPosT(0.0);
}

void PlanManager::reparamBspline(UniformBspline &bspline,
                                 vector<Eigen::Vector3d> &terminal_constraint,
                                 double ratio,
                                 Eigen::MatrixXd &ctrl_pts,
                                 double &dt)
{
  // MZY 为啥不是-6
  // m=n+p+1 => m-6=n+1+p-6=(n+1)-3
  int seg_num = bspline.getControlPoint().cols() - 3;
  // double length = bspline.getLength(0.1);
  // int seg_num = ceil(length / pp_.ctrl_pt_dist);

  bspline.lengthenTime(ratio);
  double duration = bspline.getTimeSum();
  dt = duration / double(seg_num);

  vector<Eigen::Vector3d> points;
  for (double time = 0.0; time <= duration + 1e-4; time += dt)
  {
    points.push_back(bspline.getPosT(time));
  }
  UniformBspline::BsplineFitting(dt, points, terminal_constraint, ctrl_pts);
}

// TODO 先不考虑相机转动的速度，仅考虑当前一个时刻内的
// TODO 考虑相机最大转动范围，在poscmd2odom里面设置了，之后做成一个参数,fov同理
//      FOV = 2 * atan2(0.5 * image_width, focal_length)
// TODO 先不考虑当前相机的位置
// TODO 先不考虑相机的边缘效应
// TODO 考虑相机视野规划的时候，没有把障碍物遮挡考虑进去
/**
 * @brief plan the yaw when no traj
 */
bool PlanManager::planYaw(Eigen::Vector3d pos,
                          Eigen::Quaterniond quad)
{
  /* update entropy map for dynamic obstacles */
  updateDynamicEntropyMap();

  Eigen::Vector3d rot_x = quad.toRotationMatrix().block(0, 0, 3, 1);
  double yaw_uav_now = atan2(rot_x(1), rot_x(0));
  double theta_start = yaw_uav_now - py_.max_yaw - py_.fov_2;
  std::vector<double> entropy_vec(py_.step_num, 0);

  if (!grid_map_->getEntroy(pos, theta_start, py_.delta,
                            py_.step_num, entropy_vec))
    return false;

  std::vector<int> maxIndex;
  if (!findMaxIndex(entropy_vec, py_.fov_num, 1, maxIndex))
    return false;

  // record yaw data
  size_t n = maxIndex.size();
  yaw_data_.start_time = ros::Time::now();
  yaw_data_.duration = py_.gaze_time;
  yaw_data_.yaw.assign(n + 1, 0.0);
  yaw_data_.time.assign(n + 1, 0.0);
  yaw_data_.yaw[0] = grid_map_->getCamYaw() - yaw_uav_now;
  angleLimite(yaw_data_.yaw[0]);

  for (size_t i = 0; i < maxIndex.size(); i++)
  {
    yaw_data_.yaw[i + 1] = (maxIndex[i] + py_.fov_num / 2.0) * py_.delta +
                           theta_start - yaw_uav_now;
    yaw_data_.time[i + 1] = yaw_data_.time[i] +
                            abs(yaw_data_.yaw[i + 1] - yaw_data_.yaw[i]) / py_.max_yaw_rate;
    if (i != 0) yaw_data_.time[i + 1] += py_.gaze_time;
  }

  // vis
  std::vector<Eigen::Vector3d> pts;
  for (size_t i = 0; i < n; i++)
  {
    pts.push_back(pos);
    Eigen::Vector3d dir(cos(yaw_uav_now + yaw_data_.yaw[i + 1]),
                        sin(yaw_uav_now + yaw_data_.yaw[i + 1]), 0);
    pts.push_back(pos + dir * 0.5);
  }
  visualization_->displayYawArrowList(pts, 0.05, {1, 0, 0, 1}, 10);

  return true;
}

/**
 * @brief plan the yaw with traj
 */
bool PlanManager::planYawwithTraj(Eigen::Quaterniond quad)
{
  /* update traj infomation */
  updateTrajEntropyMap();
  /* update entropy map for dynamic obstacles */
  updateDynamicEntropyMap();

  /* get entropy */
  std::vector<double> entropy_vec;
  std::vector<double> angle_vec;
  std::vector<double> entropy_max;
  std::vector<double> angle_max;
  std::vector<double> entropy_key;
  std::vector<double> angle_key;
  std::vector<std::vector<double>> entropy_vecs;
  std::vector<std::vector<double>> angel_vecs;
  std::vector<int> maxIndex;
  std::vector<double> times;
  std::vector<double> times_sample;

  ros::Time t_now_ros = ros::Time::now();
  double t_now = (t_now_ros - local_data_.start_time_).toSec();
  Eigen::Vector3d rot_x = quad.toRotationMatrix().block(0, 0, 3, 1);
  double yaw_uav_now = atan2(rot_x(1), rot_x(0));
  double yaw_now = grid_map_->getCamYaw() - yaw_uav_now; // relative yaw
  angleLimite(yaw_now);

  Eigen::Vector3d pos, vel;
  double yaw_uav;
  double t_forward = t_now + py_.forward_time;

  if (!getEntropyatTimeT(t_forward, pos, vel, yaw_uav, angle_vec, entropy_vec))
    return false;

  filterAngle(angle_vec, entropy_vec, t_forward, t_now, yaw_now);

  maxIndex = topfinder_.findTopN(entropy_vec, 3, (int)floor(py_.fov_num / 4));

  entropy_max.resize(maxIndex.size());
  angle_max.resize(maxIndex.size());
  for (size_t i = 0; i < maxIndex.size(); i++)
  {
    entropy_max[i] = entropy_vec[maxIndex[i]];
    angle_max[i] = angle_vec[maxIndex[i]];
  }
  entropy_vecs.push_back(entropy_max);
  angel_vecs.push_back(angle_max);
  times.push_back(t_forward);

  /* Get the key view point */
  double t_key;
  bool has_key = false;
  if (findKeyViewPoint(pos, t_forward, t_key)) {
    pos = local_data_.traj_pos_.getPosT(t_key);
    // record key view point for optimization
    bspline_optimizer_->setKeyViewPoint(
        pos, local_data_.start_time_ + ros::Duration(t_key));
    // cout << GREEN_B << "pos_key: " << pos.head(2).transpose() << RESET << endl;
    // cout << GREEN_B << "delta time: " << t_key - t_forward << RESET << endl;
    visualization_->displayInterPoint({pos.head(2)}, 0.4,
                                      Eigen::Vector4d(1, 0, 0, 1), 100);

    if (t_key - t_forward < py_.gaze_time * 1.5) {
      cout << RED_B << "t_key - t_forward < py_.gaze_time * 1.5 " << RESET << endl;
    }
    else {
      has_key = getEntropyatTimeT(
          t_key, pos, vel, yaw_uav, angle_vec, entropy_vec, true, t_now, yaw_now);
    }
  }

  /* fill some point along the traj */
  if (has_key) { // have found key view point
    // cout << GREEN_B << "Key point plan success!" << RESET << endl;
    maxIndex = topfinder_.findTopN(entropy_vec, 3, (int)floor(py_.fov_num / 2));

    entropy_key.resize(maxIndex.size());
    angle_key.resize(maxIndex.size());
    for (size_t i = 0; i < maxIndex.size(); i++)
    {
      entropy_key[i] = py_.key_yaw_gain * entropy_vec[maxIndex[i]];
      angle_key[i] = angle_vec[maxIndex[i]];
    }

    int num = (int)floor((t_key - t_forward) / py_.interval_time) + 1;
    num = max(2, min(num, py_.max_sample_num));
    if (num > 2) {
      std::vector<double> t_inner = getUniformSamples(t_forward, t_key, num);
      times_sample.insert(
          times_sample.end(), t_inner.begin() + 1, t_inner.end() - 1);
    }
    int num_left = py_.max_sample_num - num;
    if (num_left > 0) {
      if (num_left * py_.interval_time + t_key < local_data_.duration_) {
        for (int i = 1; i <= num_left; i++) {
          times_sample.push_back(t_key + i * py_.interval_time);
        }
      }
      else {
        num_left =
            (int)floor((local_data_.duration_ - t_key) / py_.interval_time);
        if (num_left > 0) {
          std::vector<double> t_inner =
              getUniformSamples(t_key, local_data_.duration_, num_left + 1);
          times_sample.insert(times_sample.end(), t_inner.begin() + 1, t_inner.end());
        }
      }
    }
  }
  else { // Dont found key view point
    int num_left = py_.max_sample_num - 1;
    if (num_left * py_.interval_time + t_forward < local_data_.duration_) {
      for (int i = 1; i <= num_left; i++) {
        times_sample.push_back(t_forward + i * py_.interval_time);
      }
    }
    else {
      num_left =
          (int)floor((local_data_.duration_ - t_forward) / py_.interval_time);
      if (num_left > 0) {
        std::vector<double> t_inner =
            getUniformSamples(t_forward, local_data_.duration_, num_left + 1);
        times_sample.insert(times_sample.end(), t_inner.begin() + 1, t_inner.end());
      }
    }
  }

  for (int i = 0; i < (int)times_sample.size(); i++) {
    if (getEntropyatTimeT(times_sample[i], pos, vel, yaw_uav, angle_vec,
                          entropy_vec, true, t_now, yaw_now)) {
      maxIndex = topfinder_.findTopN(entropy_vec, 3, (int)floor(py_.fov_num / 2));
      entropy_max.resize(maxIndex.size());
      angle_max.resize(maxIndex.size());
      for (size_t j = 0; j < maxIndex.size(); j++)
      {
        entropy_max[j] = entropy_vec[maxIndex[j]];
        angle_max[j] = angle_vec[maxIndex[j]];
      }
      entropy_vecs.push_back(entropy_max);
      angel_vecs.push_back(angle_max);
      times.push_back(times_sample[i]);
    }
  }

  if (has_key) {
    size_t index = times.size();
    for (size_t i = 0; i < times.size(); i++) {
      if (times[i] > t_key) {
        index = i;
        break;
      }
    }
    entropy_vecs.insert(entropy_vecs.begin() + index, entropy_key);
    angel_vecs.insert(angel_vecs.begin() + index, angle_key);
    times.insert(times.begin() + index, t_key);
  }

  gimOpter_.setAngleOptions(times, angel_vecs, entropy_vecs);
  if (!gimOpter_.optimize(t_now, yaw_now)) {
    cout << "t_now: " << t_now << endl;
    cout << "yaw: " << yaw_now << endl;
    cout << "local_data_.duration_: " << local_data_.duration_ << endl;

    // cout times
    cout << "times: {";
    for (size_t i = 0; i < times.size(); i++) {
      cout << times[i];
      if (i != times.size() - 1) cout << ", ";
    }
    cout << "}" << endl;

    // cout entropy
    cout << "entropy: " << endl;
    for (size_t i = 0; i < entropy_vecs.size(); i++)
    {
      cout << "{";
      for (size_t j = 0; j < entropy_vecs[i].size(); j++) {
        cout << entropy_vecs[i][j];
        if (j != entropy_vecs[i].size() - 1) cout << ", ";
      }
      cout << "}";
      if (i != entropy_vecs.size() - 1) cout << ", ";
      cout << endl;
    }

    // cout angle
    cout << "angle: " << endl;
    for (size_t i = 0; i < angel_vecs.size(); i++)
    {
      cout << "{";
      for (size_t j = 0; j < angel_vecs[i].size(); j++) {
        cout << angel_vecs[i][j];
        if (j != angel_vecs[i].size() - 1) cout << ", ";
      }
      cout << "}";
      if (i != angel_vecs.size() - 1) cout << ", ";
      cout << endl;
    }
    cout << RED_B << "gimOpter optimize failed" << RESET << endl;
    return false;
  }

  // cout << GREEN_B << "gimOpter optimize success" << RESET << endl;

  // record yaw data
  std::vector<std::pair<double, double>> ret = gimOpter_.getOptimalSequence();

  size_t n = ret.size();
  yaw_data_.start_time = t_now_ros;
  yaw_data_.duration = py_.gaze_time;
  yaw_data_.yaw.assign(n, 0.0);
  yaw_data_.time.assign(n, 0.0);

  for (size_t i = 0; i < n; i++)
  {
    yaw_data_.time[i] = ret[i].first;
    yaw_data_.yaw[i] = ret[i].second;
    // cout << "Time " << ret[i].first << ": Angle " << ret[i].second << endl;
  }

  // vis
  std::vector<Eigen::Vector3d> pts;
  for (size_t i = 0; i < n; i++)
  {
    Eigen::Vector3d pos = local_data_.traj_pos_.getPosT(ret[i].first);
    pts.push_back(pos);
    Eigen::Vector3d vel = local_data_.traj_vel_.getPosT(ret[i].first);
    double yaw_uav = atan2(vel(1), vel(0));
    Eigen::Vector3d dir(cos(yaw_uav + ret[i].second), sin(yaw_uav + ret[i].second), 0);
    pts.push_back(pos + dir * 0.5);
  }
  visualization_->displayYawArrowList(pts, 0.05, {1, 0, 0, 1}, 10);

  return true;
}

/**
 * @brief
 *
 * @param t time relative to the start time of the traj
 * @param pos
 * @param vel
 * @param yaw_uav
 * @param angle_vec relative angle in uav frame
 * @param entropy_vec
 * @param isLimit whether limit the angle to the last time
 * @param t_last
 * @param yaw_last relative yaw of the last time
 * @return true
 * @return false
 */
bool PlanManager::getEntropyatTimeT(double t,
                                    Eigen::Vector3d &pos, Eigen::Vector3d &vel,
                                    double &yaw_uav,
                                    std::vector<double> &angle_vec,
                                    std::vector<double> &entropy_vec,
                                    bool isLimit, double t_last, double yaw_last)
{
  if (t > local_data_.duration_ + 1e-4) {
    cout << BLUE << "t is out of the traj duration, Dont care"
         << RESET << endl;
    return false;
  }

  pos = local_data_.traj_pos_.getPosT(t);
  vel = local_data_.traj_vel_.getPosT(t);
  yaw_uav = atan2(vel(1), vel(0));
  std::vector<double> angles; // absolute angle
  std::vector<double> entropys;

  /* Sample */
  /* whether limit the angle to the last time */
  if (isLimit) {
    if (t_last >= t) {
      cout << RED_B << "[getEntropy]t_last is bigger than t" << RESET << endl;
      return false;
    }

    double delta_max = py_.max_yaw_rate * (t - t_last) * py_.runtimes_ratio;
    double yaw_l = yaw_uav - py_.fov_2 + max(yaw_last - delta_max, -py_.max_yaw);
    double yaw_h = yaw_uav + py_.fov_2 + min(yaw_last + delta_max, py_.max_yaw);
    int n = ceil(
        py_.step_num * abs(yaw_h - yaw_l) / (2 * (py_.max_yaw + py_.fov_2)));
    angles = getUniformSamples(yaw_l, yaw_h, n);
  }
  else {
    angles = getUniformSamples(yaw_uav - py_.max_yaw - py_.fov_2,
                               yaw_uav + py_.max_yaw + py_.fov_2,
                               py_.step_num);
  }

  /* get entropy */
  if (!grid_map_->getEntroy(pos, vel,
                            local_data_.start_time_ + ros::Duration(t),
                            angles, entropys))
    return false;

  /* sum up the entropy and fill the angle_vec and entropy_vec */
  int n = (int)entropys.size() - py_.fov_num + 1;
  if (n < 1) {
    n = 1;
    cout << RED_B << "[getEntropyatTimeT] n<1, BUG" << RESET << endl;
  }

  angle_vec.resize(n, 0);
  entropy_vec.resize(n, 0);

  double currentSum = 0.0;
  for (int i = 0; i < py_.fov_num; ++i)
  {
    currentSum += entropys[i];
  }

  // relative angle in uav frame
  angle_vec[0] = (angles[0] + angles[py_.fov_num - 1]) / 2 - yaw_uav;
  entropy_vec[0] = currentSum;

  for (int i = py_.fov_num; i < (int)entropys.size(); ++i)
  {
    // relative angle in uav frame
    angle_vec[i - py_.fov_num + 1] =
        (angles[i] + angles[i - py_.fov_num + 1]) / 2 - yaw_uav;
    currentSum = currentSum - entropys[i - py_.fov_num] + entropys[i];
    entropy_vec[i - py_.fov_num + 1] = currentSum;
  }

  return true;
}

/**
 * @brief filter the angle and entropy out of the range
 *
 * @param angle_vec relative angle in uav frame
 * @param yaw_last relative angle in uav frame
 */
void PlanManager::filterAngle(std::vector<double> &angle_vec,
                              std::vector<double> &entropy_vec,
                              double t, double t_last, double yaw_last)
{
  if (t_last >= t) {
    cout << RED_B << "[filterAngle]t_last is bigger than t" << RESET << endl;
    return;
  }

  size_t n = angle_vec.size();

  if (n == 0) {
    cout << RED_B << "angle_vec is empty" << RESET << endl;
    return;
  }
  if (n != entropy_vec.size()) {
    cout << RED_B << "angle_vec and entropy_vec have different size"
         << RESET << endl;
    return;
  }

  double delta_max = py_.max_yaw_rate * (t - t_last);
  double yaw_l = yaw_last - delta_max;
  double yaw_h = yaw_last + delta_max;

  for (int i = (int)n - 1; i >= 0; --i) {
    if (angle_vec[i] < yaw_l || angle_vec[i] > yaw_h) {
      angle_vec.erase(angle_vec.begin() + i);
      entropy_vec.erase(entropy_vec.begin() + i);
    }
  }
}

void PlanManager::updateTrajEntropyMap()
{
  // reset the traj map
  grid_map_->resetTrajMap();

  constexpr double time_step = 0.2; // TODO 更改时间间隔
  double t_now_gl = ros::Time::now().toSec();
  double t_start = t_now_gl - local_data_.start_time_.toSec();
  double dura = local_data_.duration_ * 2 / 3 - t_start; // TODO update2/3
  std::vector<Eigen::Vector3d> pos_vec;
  std::vector<double> t_vec;

  for (double t = time_step; t < dura; t += time_step)
  {
    t_vec.push_back(t);
    pos_vec.push_back(local_data_.traj_pos_.getPosT(t + t_start));
  }

  grid_map_->updateTrajMap(pos_vec, t_vec, dura, t_now_gl);
  grid_map_->publishTrajMap();
}

void PlanManager::updateDynamicEntropyMap()
{
  if (!have_reset_dynamic_map_) {
    grid_map_->resetDynamicObs();
    grid_map_->publishDynamicObs(true);
    have_reset_dynamic_map_ = true;
  }

  std::vector<std::vector<Eigen::Vector3d>> pos_obs;
  std::vector<std::vector<double>> R2;
  std::vector<std::vector<double>> times;
  std::vector<bool> is_dangers; // vector<bool> is stupid

  double t_now = ros::Time::now().toSec();
  if (predictor_->getPredictPos(pos_obs, R2, times, is_dangers)) {
    grid_map_->updateDynamicObs(pos_obs, R2, times, is_dangers);
    // cout << "time: (ms)" << (ros::Time::now() - t_now_ros).toSec() * 1000 << endl;
    // 0.06 ms
    grid_map_->publishDynamicObs();
    have_reset_dynamic_map_ = false;
  }
}

/**
 * @brief find the key viewpoint along the traj
 *
 * @param pos plan position
 * @param t_cur plan time relative to the start time of the traj
 * @param t_key the time of the key viewpoint
 *              relative to the start time of the traj
 */
bool PlanManager::findKeyViewPoint(
    Eigen::Vector3d pos, double t_cur, double &t_key)
{
  // get the polygon of obstacle by ray casting
  vector<Eigen::Vector2d> list_poly;
  if (!grid_map_->pi_.getPolygon(list_poly)) {
    cout << GREEN_B << "[YAW] no obstacle polygon" << RESET << endl;
    return false;
  }

  // Visualization
  visualization_->displayPolygon(list_poly, 0.2, 100);

  double t_step = 0.2; // TODO PARAM
  double dura = local_data_.duration_;
  t_key = -1.0;
  std::vector<Eigen::Vector2d> pos_vec;
  Eigen::Vector2d pos_2;
  Eigen::Vector2d pos_2_last;
  Eigen::Vector2d pos_key;

  double t_sample = t_cur + 0.05;
  pos_2 = pos_2_last = local_data_.traj_pos_.getPosT(t_sample).head(2);
  pos_vec.push_back(pos_2);

  // First point is in polygon
  if (grid_map_->pi_.isInPolygon(pos_2)) {
    bool allInPloygon = true;

    for (t_sample = t_sample + t_step; t_sample <= dura; t_sample += t_step) {
      pos_2 = local_data_.traj_pos_.getPosT(t_sample).head(2);
      pos_vec.push_back(pos_2);
      if (!grid_map_->pi_.isInPolygon(pos_2)) {
        allInPloygon = false;
        break;
      }
      pos_2_last = pos_2;
    }

    if (!allInPloygon) {
      if (grid_map_->pi_.findIntersection(pos_2_last, pos_2, pos_key)) {
        // cout << "list_poly: [";
        // for (size_t i = 0; i < list_poly.size(); i++) {
        //   cout << "[" << list_poly[i](0) << "," << list_poly[i](1) << "],";
        // }
        // cout << "]" << endl;

        // cout << "pos_vec: [";
        // for (size_t i = 0; i < pos_vec.size(); i++) {
        //   cout << "[" << pos_vec[i](0) << "," << pos_vec[i](1) << "],";
        // }
        // cout << "]" << endl;

        t_key = t_sample -
                t_step * (pos_2 - pos_key).norm() / (pos_2 - pos_2_last).norm();

        if (t_key - t_cur > py_.gaze_time * 1.5) {
          return true;
        }
        // else keep finding
      }
      else {
        visualization_->displayInterPoint(pos_vec, 0.2, Eigen::Vector4d(1, 0, 0, 1), 100);
        cout << RED_B << "[YAW]: ERROR, findIntersection failed!!"
             << RESET << endl;
        return false;
      }
    }
    else {
      cout << GREEN_B << "[YAW]: Traj are in polygon" << RESET << endl;
      return false;
    }
  }

  // First point is not in polygon or tkey is close to t_cur
  // Find the point of polygon that has the max angle with the velocity
  //   while the distance of projection is larger than a threshold
  cout << GREEN_B << "[YAW]: First point is not in polygon"
       << " or t_key is close to t_cur" << RESET << endl;
  Eigen::Vector3d vel_start = local_data_.traj_vel_.getPosT(t_cur);
  Eigen::Vector2d v_norm = vel_start.head(2).normalized();
  constexpr double min_dis_proj_vis = 0.4; // TODO PARAM
  double proj_max = 0, proj;
  double cos_min = std::numeric_limits<double>::max();
  double cos_val;
  int index = -1;
  for (size_t i = 0; i < list_poly.size() - 2; i++) {
    Eigen::Vector2d p = list_poly[i] - pos.head(2);
    proj = v_norm.dot(p);
    if (proj < min_dis_proj_vis) continue;
    cos_val = proj / p.norm();
    if (cos_val < cos_min) {
      cos_min = cos_val;
      proj_max = proj;
      index = (int)i;
    }
  }

  if (index == -1) {
    cout << GREEN_B << "[YAW]: No point in polygon is larger than min_dis_proj_vis"
         << RESET << endl;
    return false;
  }
  else {
    Eigen::Vector2d p_poly = list_poly[index];
    Eigen::Vector2d p_now = pos.head(2);
    Eigen::Vector2d p_perp, p_perp_last;
    double val = 0, val_last = 0;

    double t_estimated = t_cur + proj_max / vel_start.head(2).norm();
    t_estimated = min(t_estimated, dura);
    int count = 0;
    while (count < 100) {
      p_perp = local_data_.traj_pos_.getPosT(t_estimated).head(2);
      val = (p_perp - p_now).dot(p_poly - p_now);

      if (t_estimated > dura) {
        t_estimated = dura;
        break;
      }
      else if (t_estimated < t_cur) {
        cout << RED_B << "[YAW]: It is impossible t_estimated < t_cur"
             << RESET << endl;
        return false;
      }
      else if (val * val_last < 0) {
        t_estimated = t_estimated + t_step * val / abs(val_last - val);
        break;
      }

      if (val < -0.01) {
        t_estimated -= t_step;
      }
      else if (val > 0.01) {
        t_estimated += t_step;
      }
      else {
        break;
      }
      val_last = val;
      count++;
    }

    if (count == 100) {
      cout << RED_B << "[YAW]: count == 100, t_step is too small" << RESET << endl;
      return false;
    }

    t_key = t_estimated;
    return true;
  }

  cout << RED_B << "[YAW]: The end of findKeyViewPoint, impossible!" << RESET << endl;
  return false;
}

/**
 * @brief
 * @param ret : the index of the max value angle, not the angle
 */
bool PlanManager::findMaxIndex(const std::vector<double> &data,
                               int length, int num,
                               std::vector<int> &ret)
{
  if (data.size() < (size_t)length)
  {
    std::cerr << "Error: Insufficient data." << std::endl;
    return false;
  }

  // calculate the sum of the specific length
  int n = (int)data.size() - length + 1;
  vector<double> sum(n, 0);

  double currentSum = 0.0;
  for (int i = 0; i < length; ++i)
  {
    currentSum += data[i];
  }
  sum[0] = currentSum;

  for (int i = length; i < (int)data.size(); ++i)
  {
    currentSum = currentSum - data[i - length] + data[i];
    sum[i - length + 1] = currentSum;
  }

  int thre = (int)floor(length / 2); // TODO: 两个最大值之间的最少index数
  ret = topfinder_.findTopN(sum, num, thre);

  return true;
}

std::vector<double> PlanManager::getUniformSamples(
    double start, double end, int n)
{
  // 如果n为0，返回空数组
  if (n == 0) return std::vector<double>();

  std::vector<double> samples(n);
  if (n == 1) {
    samples[0] = start;
    return samples;
  }

  double step = (end - start) / (n - 1);

  for (int i = 0; i < n; ++i) {
    samples[i] = start + i * step;
  }

  return samples;
}

} // namespace ego_planner
